package com.ruoyi.ai.bootstrap;

import org.springframework.ai.reader.tika.TikaDocumentReader;
import com.ruoyi.ai.service.VectorService;
import com.ruoyi.ai.splitter.LineTextSplitter;
import com.ruoyi.ai.splitter.MarkdownSplitter;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.document.Document;
import org.springframework.ai.transformer.splitter.TextSplitter;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.boot.CommandLineRunner;
import org.springframework.core.io.ClassPathResource;
import org.springframework.core.io.Resource;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.List;
import java.util.Objects;


/**
 * @author 15332
 */
@Component
public class LoadVectorStore implements CommandLineRunner {
  /**
   * 日志记录器
   */
  private static final Logger logger = LoggerFactory.getLogger(LoadVectorStore.class);

  /**
   * 向量服务接口，用于商品数据的存储和检索
   */
  @Autowired
  private VectorService vectorService;

  @Autowired
  @Qualifier("milvusVectorStore")
  private VectorStore vectorStore;


  /**
   * Spring Boot应用启动时自动执行的初始化方法
   * 检查向量数据库中是否存在商品数据，如果为空则加载预设的示例商品数据
   *
   * @param args 启动参数
   * @throws Exception 初始化过程中可能发生的异常
   */
  @Override
  public void run(String... args) throws Exception {
    logger.debug("开始初始化问题数据");
    if (args.length > 0) {
      logger.info("接收到参数: {}", args[0]);
    }

    if (Objects.requireNonNull(vectorService.searchQuestion("什么是重载，重载的理解，重载的优点，重载有什么用，重载的体系介绍")).isEmpty()) {
      logger.debug("问题数据为空，重新加载");
      List<Resource> list = List.of(
        new ClassPathResource("/doc/Java基础-语法.md")
      );
      //读取文件
      list.forEach(document -> {
        TikaDocumentReader tikaDocumentReader = new TikaDocumentReader(document);
        List<Document> documents = tikaDocumentReader.get();
        TextSplitter textSplitter = null;
        //判断Resource文档类型
        if (Objects.requireNonNull(document.getFilename()).endsWith(".md")) {
          //使用MarkdownSplitter切分
          textSplitter = new MarkdownSplitter();
        } else if (document.getFilename().endsWith(".txt")) {
          textSplitter = new LineTextSplitter();
        } else {
          //使用TokenTextSplitter切分
          textSplitter = new TokenTextSplitter();
        }
        List<Document> splitDocuments = textSplitter.split(documents.get(0));
        // 分批处理文档，每批最多10个
        List<List<Document>> batches = new ArrayList<>();
        for (int i = 0; i < splitDocuments.size(); i += 10) {
          int end = Math.min(i + 10, splitDocuments.size());
          batches.add(splitDocuments.subList(i, end));
        }

        // 逐批添加到向量存储
        for (List<Document> batch : batches) {
          vectorStore.add(batch);
        }
        logger.info("初始化了{}个示例问题数据", splitDocuments.size());
      });

    }
    logger.debug("初始化数据完毕");
  }
}
